import os
import time

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader

from models.lwCETModel import lwCET
from myutils.configUtils import get_configs
from myutils.dataset import SignalDataset
from utils import to_device, save_data

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dataset_configs, hparams_class1 = get_configs()


def loadData(data_path1):
    train_data = pd.read_csv(data_path1, header=None).values
    print(train_data.shape)
    train_data = SignalDataset(train_data)
    train_dl = DataLoader(dataset=train_data,
                          batch_size=dataset_configs.batch_size,
                          num_workers=2,
                          shuffle=True)
    return train_dl


def loadModel(model_dir):
    test_model_name = os.path.join(model_dir, "checkpoint_best.pt")
    # model1 = ecgTransForm(configs=self.dataset_configs, hparams=self.hparams)
    model1 = lwCET(configs=dataset_configs, hparams=hparams_class1, add_fea=False, ablation='C')
    chkpoint = torch.load(test_model_name, map_location=device)
    model1.load_state_dict(chkpoint['model'])
    model1 = model1.to(device)
    return model1


def evaluate(model1, dataset1):
    model1.to(device).eval()

    with torch.no_grad():
        for batches in dataset1:
            batches = to_device(batches, device)
            data = batches[0].float()
            feature1 = batches[1].float()
            labels = batches[2].long()
            data = data.unsqueeze(1)
            # predictions = model(data)
            predictions = model1(data, feature1)


if __name__ == "__main__":
    exp_log_dir = "../experiments_logs/Brn/lwect-268-B_2025_10_11_11_40/"
    model = loadModel(exp_log_dir)

    data_path = "../data/mit/alldata/data_4.csv"
    dataset = loadData(exp_log_dir)

    times = []
    for i in range(60):
        start = time.time()
        evaluate(model, dataset)
        cost_time = time.time() - start
        times.append(cost_time)

    save_data("train_time.txt", times, True)
